# -*- coding: utf-8 -*-
"""
文件名称: disease_keep_ratio_stats.py

功能概述:
    基于支持文件 (support_cleaned.json) 中的疾病标签频次, 使用“频率驱动 + 分段线性”的策略
    为每一种疾病组合分配一个保留概率 keep_ratio, 并将结果写入一个纯统计用的 JSON 文件。

    这里的 keep_ratio 含义为:
        - 对于某个疾病 key, 若 keep_ratio = 0.05,
          表示“在处理该疾病相关的样本时, 以 5% 的概率保留该样本”。
          (实际实现中是对样本做一次 [0,1) 的均匀采样, 若 r <= keep_ratio 则保留)

    支持文件格式示意:
        {
          "aptos_001639a390f0.png": {
            "image_path": ".../aptos_001639a390f0.png",
            "label": ["增殖型糖尿病视网膜病变"],
            "quality": null
          },
          "messidor_20060411_58550_0200_PP.png": {
            "image_path": ".../messidor_20060411_58550_0200_PP.png",
            "label": null,
            "quality": "图像质量不可评估..."
          },
          ...
        }

    统计逻辑:
        1) 对于每个条目:
            - 若 label 为 None 或空列表 []:
                  认为没有有效疾病标签, 不参与疾病统计。
            - 若 label 为非空 list:
                  将其中元素转为 str, 排序后用 "_" join, 得到 disease_key。
        2) 统计每个 disease_key 的频次 freq, 并基于频次为其分配 keep_ratio。
        3) 输出 JSON, 核心为一个三元组列表:
            [
              [疾病名称(字符串), 该疾病的图像数量 freq, 保留概率 keep_ratio],
              ...
            ]
           另外附带总疾病数、总带标签图像数、使用的超参数以及一个“等效删除统计”。

    当前采用的 keep_ratio 设计 (频次记为 f, 分三段 + Top3 特殊处理):

        * 记 diseases 按频次从高到低排序, 得到:
              (d1, f1), (d2, f2), (d3, f3), (d4, f4), (d5, f5), ...

          其中:
              - 前 3 个疾病 d1, d2, d3 直接固定 keep_ratio = 0.05
                (最常见的 3 类强力压缩)
              - f5 作为“高频段”的上界, 用于线性插值的最大频次

        * 基本分段规则 (针对非 Top3 疾病):

          1) 低频区: f <= 20
             -> keep_ratio = 1.0
             (20 次以内全部保留, 完全不动长尾)

          2) 中频区: 20 < f <= 1000
             -> 线性从 1.0 衰减到 0.2:

                t = (f - 20) / (1000 - 20)
                keep(f) = 1.0 - t * (1.0 - 0.2)
                         = 1.0 - t * 0.8

             即:
                f =  20  -> keep=1.0
                f = 1000 -> keep≈0.2
                中间平滑过渡

          3) 高频区: f > 1000
             -> 使用“第五大频次 f5”作为高频段的上界:
                若 f5 > 1000:
                    t = (f - 1000) / (f5 - 1000)
                    keep(f) = 0.2 - t * (0.2 - 0.05)
                             = 0.2 - t * 0.15
                    然后与 0.05 取最大值:
                        keep(f) = max(0.05, keep(f))

                若 f5 <= 1000 (极端情况, 实际上几乎只会出现在小数据集):
                    当前数据集没有真正意义上的“>1000 的高频段”,
                    则对 f>1000 的非 Top3 疾病统一采用 keep_ratio=0.2 兜底。

        * Top3 特殊处理:
            - 对于疾病 d1, d2, d3 (频次最高的三个), 无论频次是多少, 一律:
                  keep_ratio = 0.05

    等效删除的定义 (全局统计):
        - 记 disease_freq[d] 为疾病 d 的样本数, keep_ratio[d] 为其保留概率。
        - 有效保留比例:
              effective_keep_ratio =
                  sum_d ( disease_freq[d] * keep_ratio[d] ) / sum_d ( disease_freq[d] )
        - 等效删除比例:
              effective_delete_ratio = 1 - effective_keep_ratio

      直观理解:
        假设你对整库做“统一保留 p%”的抽样, 想要和现在这套“按疾病不同 p_d%”的策略总体效果一致,
        那这个等效 p 就是 effective_keep_ratio, 对应的整体删除比例就是 effective_delete_ratio。

    输出 JSON 结构示例:
        {
          "total_diseases": 389,
          "total_images_with_label": 334086,
          "params": {
            "low_full_keep_threshold": 20,
            "mid_start": 20,
            "mid_end": 1000,
            "top3_fixed_keep": 0.05
          },
          "effective_keep_ratio": 0.1732,
          "effective_delete_ratio": 0.8268,
          "items": [
            ["正常眼底", 235391, 0.05],
            ["中度糖尿病视网膜病变", 31086, 0.07],
            ...
          ]
        }

使用示例:
    python disease_keep_ratio_stats.py \
        --support-cleaned-path /path/to/cls_annotations_merged.json \
        -o /path/to/disease_keep_ratio_stats.json
"""

import argparse
import json
import logging
from collections import Counter
from pathlib import Path
from typing import Any, Dict, List, Tuple


# -----------------------------
# 全局参数 (分段规则)
# -----------------------------

# 低频保留阈值: f <= LOW_FULL_KEEP_THRESHOLD -> keep_ratio = 1.0
LOW_FULL_KEEP_THRESHOLD = 20

# 中频区间: (MID_START, MID_END] 做线性插值
MID_START = 20
MID_END = 1000

# Top3 疾病固定保留比例
TOP3_KEEP_RATIO = 0.05

# 中频段末端的目标保留比例
MID_END_KEEP_RATIO = 0.20

# 高频段最低保留比例 (Top3 之外的疾病不会低于该值)
HIGH_FREQ_MIN_KEEP_RATIO = 0.05

logger = logging.getLogger(__name__)
if not logger.handlers:
    _handler = logging.StreamHandler()
    _formatter = logging.Formatter(
        fmt="[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    _handler.setFormatter(_formatter)
    logger.addHandler(_handler)
    logger.setLevel(logging.INFO)


# -----------------------------
# 读取支持文件 & 统计疾病频次
# -----------------------------

def load_support_cleaned(support_path: Path) -> Dict[str, Any]:
    """
    读取支持 JSON 文件, 返回原始 dict。

    参数:
        support_path: 支持 JSON 文件路径

    返回:
        support_dict: 原始 JSON 解析后的 dict

    异常:
        - 文件不存在或无法读取时抛出异常
        - JSON 解析失败时抛出异常
        - 顶层结构不是 dict 时抛出异常
    """
    if not support_path.is_file():
        raise FileNotFoundError(f"support_cleaned_path 文件不存在: {support_path}")

    try:
        text = support_path.read_text(encoding="utf-8")
        data = json.loads(text)
    except Exception as e:
        raise RuntimeError(f"读取或解析支持 JSON 失败: path={support_path}, error={e}") from e

    if not isinstance(data, dict):
        raise RuntimeError(
            f"支持 JSON 顶层结构必须为 dict, 实际为: {type(data)} (path={support_path})"
        )

    return data


def build_disease_freq(
    support_dict: Dict[str, Any]
) -> Tuple[Dict[str, int], int]:
    """
    基于支持 JSON, 统计 “疾病 key” 的出现频次。

    疾病 key 的构造方式:
        - 对于 support_dict 中的每个条目:
            * key: 图像文件名, 如 "aptos_001639a390f0.png"
            * value: 记录 dict, 其中 "label" 字段为疾病标签列表
        - 对 label 字段:
            * 若 label 为 None 或空列表 []:
                - 视为没有可用于疾病统计的信息, 跳过。
            * 若 label 为非空 list:
                - 将元素转为 str, 排序后用 "_" join, 得到 disease_key。
                - disease_freq[disease_key] += 1

    返回:
        disease_freq: {disease_key: 出现次数}
        total_images_with_label: 具有非空 label 的图像总数
    """
    freq_counter: Counter = Counter()
    total_with_label = 0

    for img_name, info in support_dict.items():
        if not isinstance(info, dict):
            raise RuntimeError(
                f"support_cleaned 中 key={img_name!r} 对应的 value 非 dict 类型: {type(info)}"
            )

        labels = info.get("label", None)

        # label 为 None 或空列表 -> 不参与疾病统计
        if labels is None:
            logger.debug(
                "support_cleaned: key=%s 的 label 为 None, 跳过疾病统计",
                img_name,
            )
            continue
        if isinstance(labels, list) and len(labels) == 0:
            logger.debug(
                "support_cleaned: key=%s 的 label 为空列表, 跳过疾病统计",
                img_name,
            )
            continue

        if not isinstance(labels, list):
            raise RuntimeError(
                f"support_cleaned 中 key={img_name!r} 的 label 类型异常, "
                f"期望为 list / None, 实际为: {type(labels)}, value={labels!r}"
            )

        labels_sorted = sorted(str(x) for x in labels)
        disease_key = "_".join(labels_sorted)

        freq_counter[disease_key] += 1
        total_with_label += 1

    disease_freq: Dict[str, int] = dict(freq_counter)
    if not disease_freq:
        raise RuntimeError("在 support_cleaned 中未统计到任何疾病 key, 请检查输入文件与 label 字段")

    logger.info(
        "疾病频次统计完成: 不同疾病 key 数量=%d, 具有非空标签的图像总数=%d",
        len(disease_freq),
        total_with_label,
    )
    return disease_freq, total_with_label


# -----------------------------
# 频率驱动 keep_ratio 计算 (新规则)
# -----------------------------

def build_disease_keep_ratio(
    disease_freq: Dict[str, int],
) -> Dict[str, float]:
    """
    根据疾病在数据集中的“绝对频次”动态分配 keep_ratio, 更温和地削减高频疾病,
    避免对头部类进行过度删除, 同时基本不动长尾。

    设计 (频次记为 f, 分三段 + Top3 特殊处理):

        1) 排序:
            - 按频次从高到低排序得到列表:
                  sorted_desc = [(d1, f1), (d2, f2), (d3, f3), (d4, f4), (d5, f5), ...]

            - Top3 疾病 d1, d2, d3:
                  keep_ratio = TOP3_KEEP_RATIO (默认 0.05)

            - f5 作为高频段插值的上界:
                  若疾病数 >= 5, 则 f5 = 第五大频次
                  否则退化为 max_freq (防御性处理)

        2) 分段规则 (针对非 Top3 疾病):

            (a) 低频区: f <= 20
                -> keep_ratio = 1.0

            (b) 中频区: 20 < f <= 1000
                -> 从 1.0 线性衰减到 0.2:

                   t = (f - 20) / (1000 - 20)
                   keep(f) = 1.0 - t * (1.0 - 0.2) = 1.0 - 0.8 * t

            (c) 高频区: f > 1000

                - 若 f5 > 1000:
                    t = (f - 1000) / (f5 - 1000)
                    keep(f) = 0.2 - t * (0.2 - 0.05)
                             = 0.2 - 0.15 * t
                    最后与 0.05 取最大值:
                        keep(f) = max(0.05, keep(f))

                - 若 f5 <= 1000 (极端情况):
                    -> 非 Top3 的 f>1000 疾病一律 keep_ratio = 0.2 兜底。

    返回:
        disease_keep_ratio: {disease_key: keep_ratio ∈ (0, 1]}
    """
    if not disease_freq:
        raise RuntimeError("build_disease_keep_ratio: disease_freq 为空, 无法分配 keep_ratio")

    # 按频次从高到低排序
    sorted_desc = sorted(disease_freq.items(), key=lambda x: x[1], reverse=True)
    max_freq = sorted_desc[0][1]
    min_freq = min(disease_freq.values())

    # Top3 疾病
    top3_keys = {k for k, _ in sorted_desc[:3]}

    # 第五大频次 (若不足 5 个疾病, 退化为 max_freq)
    if len(sorted_desc) >= 5:
        f5 = float(sorted_desc[4][1])
    else:
        f5 = float(max_freq)

    logger.info(
        "疾病频次范围: min_freq=%d, max_freq=%d, 第五大频次 f5=%.0f",
        min_freq,
        max_freq,
        f5,
    )
    logger.info(
        "Top3 疾病固定 keep_ratio=%.3f, 分别为: %s",
        TOP3_KEEP_RATIO,
        ", ".join(list(top3_keys)),
    )

    disease_keep_ratio: Dict[str, float] = {}

    for d_key, freq in disease_freq.items():
        f = float(freq)

        # 1) Top3 疾病: 固定 0.05
        if d_key in top3_keys:
            keep = TOP3_KEEP_RATIO

        # 2) 低频区: f <= 20 -> 全保留
        elif f <= LOW_FULL_KEEP_THRESHOLD:
            keep = 1.0

        # 3) 中频区: (20, 1000] -> 线性从 1.0 衰减到 0.2
        elif f <= MID_END:
            # 防御: 避免除零
            if MID_END == MID_START:
                keep = MID_END_KEEP_RATIO
            else:
                t = (f - MID_START) / (MID_END - MID_START)
                keep = 1.0 - t * (1.0 - MID_END_KEEP_RATIO)
                # 理论上落在 [0.2,1.0] 内

        # 4) 高频区: f > 1000
        else:
            # 正常情况: f5 > MID_END, 在 (1000, f5] 上线性从 0.2 衰减到 0.05
            if f5 > MID_END:
                denom = (f5 - MID_END) if f5 != MID_END else 1.0
                t = (f - MID_END) / denom
                # t 理论上在 [0,1], 做个裁剪以防数值误差
                if t < 0.0:
                    t = 0.0
                elif t > 1.0:
                    t = 1.0
                keep = MID_END_KEEP_RATIO - t * (MID_END_KEEP_RATIO - HIGH_FREQ_MIN_KEEP_RATIO)
                # 与 0.05 取最大值, 避免过度压缩
                if keep < HIGH_FREQ_MIN_KEEP_RATIO:
                    keep = HIGH_FREQ_MIN_KEEP_RATIO
            else:
                # 极端情况: f5 <= 1000, 说明实际上没有真正意义上的“高频段”
                # 但当前 f > 1000 且不是 Top3, 则统一采用 0.2 兜底
                keep = MID_END_KEEP_RATIO

        # 数值安全: keep ∈ (0,1]
        if keep > 1.0:
            keep = 1.0
        if keep <= 0.0:
            keep = HIGH_FREQ_MIN_KEEP_RATIO

        disease_keep_ratio[d_key] = float(keep)

    # 简单日志: 看看 keep_ratio 大致分布 (按疾病数)
    summary_counter: Dict[float, int] = {}
    for v in disease_keep_ratio.values():
        summary_counter[v] = summary_counter.get(v, 0) + 1

    logger.info("频率驱动 keep_ratio 分布概况(按疾病数):")
    for kr, cnt in sorted(summary_counter.items(), key=lambda x: x[0], reverse=True):
        logger.info("  keep_ratio=%.3f -> disease_count=%d", kr, cnt)

    return disease_keep_ratio


# -----------------------------
# 等效删除比例统计
# -----------------------------

def compute_effective_keep_and_delete(
    disease_freq: Dict[str, int],
    disease_keep_ratio: Dict[str, float],
) -> Tuple[float, float]:
    """
    计算“等效保留比例”和“等效删除比例”。

    定义:
        total_freq = sum_d disease_freq[d]
        kept_freq_est = sum_d ( disease_freq[d] * keep_ratio[d] )

        effective_keep_ratio   = kept_freq_est / total_freq
        effective_delete_ratio = 1 - effective_keep_ratio

    直观理解:
        - effective_keep_ratio 表示: 如果对整个数据集做一个统一的 p% 抽样,
          其总体删除效果与当前“按疾病分段抽样”的策略相当, 那么这个 p 就是 effective_keep_ratio。
        - effective_delete_ratio 则是对应的等效删除比例。
    """
    if not disease_freq:
        raise RuntimeError("compute_effective_keep_and_delete: disease_freq 为空")

    total_freq = 0.0
    kept_freq_est = 0.0

    for d_key, freq in disease_freq.items():
        total_freq += float(freq)
        kr = disease_keep_ratio.get(d_key)
        if kr is None:
            raise RuntimeError(
                f"compute_effective_keep_and_delete: 疾病 key={d_key!r} 在 disease_keep_ratio 中未找到"
            )
        kept_freq_est += float(freq) * float(kr)

    if total_freq <= 0.0:
        raise RuntimeError("compute_effective_keep_and_delete: total_freq <= 0, 无法计算比例")

    effective_keep_ratio = kept_freq_est / total_freq
    effective_delete_ratio = 1.0 - effective_keep_ratio

    logger.info(
        "等效保留比例: %.4f (%.2f%%), 等效删除比例: %.4f (%.2f%%)",
        effective_keep_ratio,
        effective_keep_ratio * 100.0,
        effective_delete_ratio,
        effective_delete_ratio * 100.0,
    )

    return effective_keep_ratio, effective_delete_ratio


# -----------------------------
# 输出 JSON 组装
# -----------------------------

def build_output_json(
    disease_freq: Dict[str, int],
    total_images_with_label: int,
    disease_keep_ratio: Dict[str, float],
    effective_keep_ratio: float,
    effective_delete_ratio: float,
) -> Dict[str, Any]:
    """
    构造输出 JSON 对象。

    输出结构:
        {
          "total_diseases": ...,
          "total_images_with_label": ...,
          "params": {...},
          "effective_keep_ratio": ...,
          "effective_delete_ratio": ...,
          "items": [
            [disease_key, freq, keep_ratio],
            ...
          ]
        }
    """
    # 为了更直观, 按 freq 从高到低排序
    items: List[Tuple[str, int, float]] = []
    for d_key, freq in disease_freq.items():
        kr = disease_keep_ratio.get(d_key)
        if kr is None:
            raise RuntimeError(
                f"疾病 key={d_key!r} 在 disease_keep_ratio 中未找到, "
                "请检查 build_disease_keep_ratio 逻辑"
            )
        items.append((d_key, freq, kr))

    items_sorted = sorted(items, key=lambda x: x[1], reverse=True)

    out = {
        "total_diseases": len(disease_freq),
        "total_images_with_label": int(total_images_with_label),
        "params": {
            "low_full_keep_threshold": LOW_FULL_KEEP_THRESHOLD,
            "mid_start": MID_START,
            "mid_end": MID_END,
            "top3_fixed_keep": TOP3_KEEP_RATIO,
            "mid_end_keep_ratio": MID_END_KEEP_RATIO,
            "high_freq_min_keep_ratio": HIGH_FREQ_MIN_KEEP_RATIO,
        },
        "effective_keep_ratio": float(effective_keep_ratio),
        "effective_delete_ratio": float(effective_delete_ratio),
        # 三元组列表: [疾病名称, 频次(图像数量), 保留概率]
        "items": [
            [d_key, int(freq), float(kr)]
            for d_key, freq, kr in items_sorted
        ],
    }
    return out


# -----------------------------
# 主流程
# -----------------------------

def run_disease_keep_ratio_stats(
    support_path: Path,
    out_path: Path,
) -> None:
    """
    主流程:
        1. 读取支持 JSON
        2. 统计疾病 key 频次
        3. 基于频次计算 keep_ratio (新规则)
        4. 计算等效保留比例 & 等效删除比例
        5. 组装为纯统计 JSON 并写出
    """
    support_dict = load_support_cleaned(support_path)
    disease_freq, total_images_with_label = build_disease_freq(support_dict)
    disease_keep_ratio = build_disease_keep_ratio(disease_freq)
    effective_keep_ratio, effective_delete_ratio = compute_effective_keep_and_delete(
        disease_freq, disease_keep_ratio
    )

    out_obj = build_output_json(
        disease_freq=disease_freq,
        total_images_with_label=total_images_with_label,
        disease_keep_ratio=disease_keep_ratio,
        effective_keep_ratio=effective_keep_ratio,
        effective_delete_ratio=effective_delete_ratio,
    )

    try:
        out_path.parent.mkdir(parents=True, exist_ok=True)
        out_path.write_text(
            json.dumps(out_obj, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        logger.info("疾病频率 + 保留概率 + 等效删除统计结果已写入: %s", out_path)
    except Exception as e:
        raise RuntimeError(
            f"写出统计 JSON 失败: path={out_path}, error={e}"
        ) from e


# -----------------------------
# CLI
# -----------------------------

def parse_args() -> argparse.Namespace:
    """
    命令行参数解析。

    参数:
        --support-cleaned-path, -s  支持 JSON 文件路径
        --output, -o                输出 JSON 文件路径
    """
    parser = argparse.ArgumentParser(
        description="基于疾病频次的保留概率(keep_ratio)统计工具 (带 Top3 固定 + 分段线性策略 + 等效删除统计)",
    )
    parser.add_argument(
        "--support-cleaned-path",
        "-s",
        type=str,
        required=True,
        help="支持 JSON 文件路径 (key 为图像文件名, 含 label 字段)",
    )
    parser.add_argument(
        "--output",
        "-o",
        type=str,
        required=True,
        help="输出统计结果 JSON 文件路径",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    support_path = Path(args.support_cleaned_path)
    out_path = Path(args.output)

    run_disease_keep_ratio_stats(
        support_path=support_path,
        out_path=out_path,
    )


if __name__ == "__main__":
    main()
